from elasticsearch_dsl import Search, Q, connections
from pymilvus import Collection, connections as milvus_connections
import numpy as np
from rag_test.pingan.pingan_text_embedding import text_embeddings


# ... existing code ...
def es_keyword_search(query, index_name="insurance_manual_es_01", top_k=5):
    """
    使用Elasticsearch进行关键词搜索

    参数:
        query: 搜索关键词
        index_name: Elasticsearch索引名称
        top_k: 返回结果数量

    返回:
        包含搜索结果的列表，每个结果包含文本、页码、得分和来源
    """
    connections.create_connection(hosts=["http://localhost:9200"])
    milvus_connections.connect(host="localhost", port="19530")
    if not connections.get_connection():
        connections.create_connection(hosts=["localhost:9200"])

    # 构建并执行Elasticsearch搜索查询
    s = Search(index=index_name).using(connections.get_connection()).query(
        Q("multi_match", query=query, fields=["text"], fuzziness="AUTO")
    )[:top_k]

    response = s.execute()
    results = []
    # 处理搜索结果
    for hit in response:
        results.append({
            "text": hit.text,
            "page_num": hit.page_num,
            "score": hit.meta.score,
            "source": "es"
        })
    print("ES搜索结果：", results)
    return results


# ... existing code ...


def milvus_semantic_search(query, collection_name="insurance_manual_01", top_k=5):
    """
    使用Milvus进行语义搜索

    参数:
        query: 搜索查询语句
        collection_name: Milvus集合名称
        top_k: 返回结果数量

    返回:
        包含语义搜索结果的列表，每个结果包含文本、页码、得分和来源
    """
    connections.create_connection(hosts=["http://localhost:9200"])
    milvus_connections.connect(host="localhost", port="19530")
    if not milvus_connections.has_connection("default"):
        milvus_connections.connect(
            alias="default",
            host="localhost",
            port="19530"
        )

    # 获取查询语句的嵌入向量
    query_embedding = text_embeddings(query)[0]
    collection = Collection(collection_name)

    # 检查集合是否已加载
    try:
        loaded = collection.is_loaded
    except AttributeError:
        loaded = collection.has_index()

    if not loaded:
        collection.load()

    # 执行向量搜索
    search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
    results = collection.search(
        data=[query_embedding],
        anns_field="embedding",
        param=search_params,
        limit=top_k,
        output_fields=["text", "page_num"]
    )

    # 处理搜索结果
    semantic_results = []
    for hits in results:
        for hit in hits:
            score = 1 / (1 + hit.distance)
            semantic_results.append({
                "text": hit.entity.get("text"),
                "page_num": hit.entity.get("page_num"),
                "score": score,
                "source": "milvus"
            })
    print("Milvus搜索结果：", semantic_results)
    return semantic_results





def hybrid_search(query, es_weight=0.2, milvus_weight=0.8, top_k=5):
    # 执行关键词搜索和语义搜索
    es_results = es_keyword_search(query, top_k=top_k)
    milvus_results = milvus_semantic_search(query, top_k=top_k)

    # 合并关键词搜索和语义搜索的结果，去重并保留高分结果
    merged = {}
    for res in es_results + milvus_results:
        key = res["text"]
        if key not in merged:
            merged[key] = res
        else:
            if res["score"] > merged[key]["score"]:
                merged[key] = res
    # 计算混合得分：对不同来源的结果使用不同的权重
    for res in merged.values():
        if res["source"] == "es":
            res["hybrid_score"] = (res["score"] / es_results[0]["score"] if es_results else 0) * es_weight
        else:
            res["hybrid_score"] = res["score"] * milvus_weight

    return sorted(merged.values(), key=lambda x: x["hybrid_score"], reverse=True)[:top_k]


if __name__ == "__main__":
    connections.create_connection(hosts=["http://localhost:9200"])
    milvus_connections.connect(host="localhost", port="19530")

    user_query = "平安福2024版有什么变化？"
    results = hybrid_search(user_query)

    print(f"查询: {user_query}\n")
    for i, res in enumerate(results, 1):
        print(f"结果{i}（页码：{res['page_num']}，分数：{res['hybrid_score']:.2f}）:")
        print(res["text"] + "\n")
